import numpy as np
from numpy import pi, sin, cos
import matplotlib.pyplot as plt

# 如果导入报错了，请手动导入
from min_ROC.goat import smgoat

I, sx, sy, sz = np.identity(2, dtype=complex), np.array([[0, 1], [1, 0]], dtype=complex), np.array(
	[[0, -1.j], [1.j, 0.]]), np.diag([1., -1]).astype(complex)

max_ple, max_ore = 0.45, 0.18  # 45%, 18%

"""Ucps family"""
# U5c(Knill) pulse as initial guess
# phis0 = np.array([1 / 6, 0, 1 / 2, 0, 1 / 6]) * pi

# U7a pulse as initial guess
# phis0 = np.array([0, 11, 10, 17, 10, 11, 0]) * pi / 12  # U7a
# U7b pulse as initial guess
# phis0 = np.array([0, 23, 10, 5, 10, 23, 0]) * pi / 12  # U7b

# U9a pulse as initial guess
phis0 = np.array([0, 0.635, 1.35, 0.553, 0.297, 0.553, 1.35, 0.635, 0]) * pi
# U9b pulse as initial guess
# phis0 = np.array([0, 1.635, 1.35, 1.553, 0.297, 1.553, 1.35, 1.635, 0]) * pi

# U13a pulse as initial guess
# phis0 = np.array([0, 9, 42, 11, 8, 37, 2, 37, 8, 11, 42, 9, 0]) * pi / 24  # U13a
# U13b pulse as initial guess
# phis0 = np.array([0, 33, 42, 35, 8, 13, 2, 13, 8, 35, 42, 33, 0]) * pi / 24  # U13b

thetas0 = np.ones(len(phis0)) * pi

"""CCCPs family"""

# s, k = np.arccos(-0.125), np.arcsin(sin(pi / 4) / 2.0)
# thetas0 = [pi, 2 * pi, pi, 2 * pi + pi / 4 - k, 2 * pi - 2 * k, pi / 4 - k]
# phis0 = [s, 3 * s, s, 0, pi, 0.]
# thetas0 = np.array([2 * pi + pi / 4 - k, 2 * pi - 2 * k, pi / 4 - k, 2 * pi, 2 * pi])
# phis0 = np.array([0, pi, 0., -s, s])
def Rx(theta: float, phi: float):
	return cos(theta / 2) * I - 1.j * sin(theta / 2) * (cos(phi) * sx + sin(phi) * sy)

def get_zphase(thetas, phis, UF):
	Uf = I.copy()
	for theta, phi in zip(thetas, phis):
		h_theta = theta / 2.0
		Uf = (cos(h_theta) * I - 1.j * sin(h_theta) * (cos(phi) * sx + sin(phi) * sy)) @ Uf
	Utmp = UF.T.conj() @ Uf
	z_phase = np.angle(Utmp[0, 0]) - np.angle(Utmp[1, 1])
	Uz = cos(z_phase / 2) * I - 1.j * sin(z_phase / 2) * sz
	return z_phase, Uf @ Uz

thetaT, phiT = pi, 0.
UT = Rx(thetaT, phiT)
angle_mod = False  # True
w_max = 0.01 * (2 * pi)  # f=0.01·(2π)MHz, 1MHz <===> 1μs
sample_num = 45
result = smgoat(
	thetas0=thetas0,
	phis0=phis0,
	UT=UT,  # target unitary
	quad_level=[5, 5],
	max_ple=max_ple,
	max_ore=max_ore,  # in unit of w_max
	angle_mod=angle_mod,
	quad_type='uniform',
	infide_type='pop2 zphase',

	iters=2000,
	options={
		'maxcor': 25,
	},
)

print(result)
params = result.x

thetas_res = params[:len(thetas0)] if angle_mod else thetas0
phis_res = params[len(thetas0):] if angle_mod else params
# zphase, gate = get_zphase(thetas_res, phis_res, -1.j * sx)

from functools import partial
from ..utils import compare_robust_landscape, composite_pulses, Square, repr_composite_pulse, _angle2str

print(f"thetas = {np.array2string(thetas_res, precision=5, separator=', ', suppress_small=True)}")
print(f"phis = {np.array2string(phis_res, precision=5, separator=', ', suppress_small=True)}")
print(
	f"thetas = np.array({np.array2string(thetas_res / pi, precision=5, separator=', ', suppress_small=True)}) * np.pi")
print(f"phis = np.array({np.array2string(phis_res / pi, precision=5, separator=', ', suppress_small=True)}) * np.pi\n")
# print(f"zphase = {zphase / pi:.4} π\n gate = {np.array2string(gate, precision=4, separator=', ')}")
print(repr_composite_pulse(thetas_res, phis_res))

ple_arr, ore_arr = np.linspace(-max_ple, max_ple, sample_num), np.linspace(-max_ore, max_ore, sample_num) * w_max

test_pulse_S = {
	'benchmark': partial(composite_pulses(thetas0, phis0), Square),
	'test': partial(composite_pulses(thetas_res, phis_res), Square),
}
fig = plt.figure(figsize=(2.5 * len(test_pulse_S), 3.0), layout="constrained")
compare_robust_landscape(thetaT, phiT, w_max, ple_arr, ore_arr, Square, test_pulse_S, fig=fig,
                         infide_type='corr2 upto z phase', )

fig.suptitle(
	r'Robust Landscape of $\beta_z$,$\beta_\Omega$ when $\theta=%s$,$\phi=%s$, max_ple=$\pm$%s%% max_ore=$\pm$%s%%' % (
		_angle2str(thetaT), _angle2str(phiT), max_ple * 100, max_ore * 100))
plt.show()
